import os
import kp
import json
from tpuKneron import *
from utils.ExampleHelper import get_device_usb_speed_by_port_id


class TPU:
    def __init__(self):
        # 获取当前脚本的绝对路径
        self.pwd = os.path.dirname(os.path.abspath(__file__))
        # USB端口号、固件路径
        self.usb_port_id = 0
        self.scpu_fw_path = os.path.join(self.pwd, '../../resource/KL520/fw_scpu.bin')
        self.ncpu_fw_path = os.path.join(self.pwd, '../../resource/KL520/fw_ncpu.bin')

    # TPU检测
    def has_tpu_kneron(self):
        device_list = kp.core.scan_devices()
        print(f'Number of Kneron devices found: {device_list.device_descriptor_number}')
        return device_list.device_descriptor_number

    # 识别链接TPU
    def link(self):
        try:
            res = self.has_tpu_kneron()
            print("TPU识别成功!") if res != 0 else print("TPU无法识别, res:", res)
            try:
                if kp.UsbSpeed.KP_USB_SPEED_HIGH != get_device_usb_speed_by_port_id(usb_port_id=self.usb_port_id):
                    print('\033[91m' + '[Warning] Device is not run at high speed.' + '\033[0m')
            except Exception as exception:
                print('Error: check device USB speed fail, port ID = \'{}\', error msg: [{}]'.format(self.usb_port_id,
                                                                                                    str(exception)))
                return False
            try:
                print('[Connect Device]')
                self.device_group = kp.core.connect_devices(usb_port_ids=[self.usb_port_id])
            except kp.ApiKPException as exception:
                print('Error: connect device fail, port ID = \'{}\', error msg: [{}]'.format(self.usb_port_id,
                                                                                            str(exception)))
                return False

            kp.core.set_timeout(device_group=self.device_group, milliseconds=5000)

            try:
                print('[Upload Firmware]')
                kp.core.load_firmware_from_file(device_group=self.device_group,
                                                scpu_fw_path=self.scpu_fw_path,
                                                ncpu_fw_path=self.ncpu_fw_path)
            except kp.ApiKPException as exception:
                print('Error: upload firmware failed, error = \'{}\''.format(str(exception)))
                return False
            print("TPU链接成功!")

        except Exception as e:
            print(f"TPU初始化失败, error:{str(e)}")
        return True

    # TPU模型初始化
    def model_init(self, model_path, host=False):
        try:
            if host==False:
                self.model_nef_descriptor = kp.core.load_model_from_flash(device_group=self.device_group)
            else:
                self.model_nef_descriptor = kp.core.load_model_from_file(device_group=self.device_group,
                                                                file_path=f"{model_path}/models_520.nef")
        except kp.ApiKPException as exception:
            print('Error: upload model failed, error = \'{}\''.format(str(exception)))
            exit(0)

    # 释放TPU
    def release(self):
        try:
            kp.core.disconnect_devices(self.device_group)
            print('TPU释放成功')
        except kp.ApiKPException as exception:
            print('Error: release device fail, error = \'{}\''.format(str(exception)))
            exit(0)

    # 重启TPU
    def reboot(self):
        print("rebooting tpu...")

    # 获取TPU内部的配置文件
    def get_config(self, app_name, dst_path):
        self.release()
        tpu = tpuKneron()
        code = tpu.init()
        data = None
        if code == 0:
            data = tpu.get_model_config()
            data = json.loads(data)
            print(dst_path)
            os.makedirs(dst_path, exist_ok=True)
            # 写入到 model.json 文件
            with open(f'{dst_path}/model.conf', 'w') as json_file:
                json.dump(data, json_file, indent=4)
                return True
        else:
            return False